from typing import List

from pydantic import BaseModel


class Service(BaseModel):
    name: str
    host_port: int
    envd_port: int


class GitConfig(BaseModel):
    name: str
    email: str
    editor: str = "vim"


class OpenCVLib(BaseModel):
    pass


class GolangLib(BaseModel):
    version: str = "1.18.8"


class TensorboardLib(BaseModel):
    envd_port: int = 6006
    envd_dir: str = "/home/envd/logs",
    host_port: int = 0,
    host_dir: str = "/tmp"


class TensorRTLib(BaseModel):
    os: str = "20.04"
    cuda: str = "11.6.2"
    trt: str = "8.4.3.1"


class JupyterLabLib(BaseModel):
    envd_port: int = 8888
    host_port: int = 0
    token: str = "''"


class CudaConfig(BaseModel):
    version: str = "11.6.2"
    cudnn: str = "8"


class EnvdBuild(BaseModel):
    os: str = "ubuntu20.04"
    language: str = "python3"
    shell: str = "zsh"
    pip_index: str | None = "https://pypi.tuna.tsinghua.edu.cn/simple"
    conda_channel: str | None = None
    docker_mirror: str | None = "https://docker.mirrors.sjtug.sjtu.edu.cn"
    python_pkgs: List[str] | None = []
    vscode_exts: List[str] | None = [
        "ms-python.python",
        "ms-toolsai.jupyter",
    ]
    expose_service: List[Service] | None = None
    daemon_commands: List[str] | None = None
    cuda: CudaConfig | None = None,
    git: GitConfig | None = None,
    libs: dict | None = None


class EnvdInDB(BaseModel):
    project_name: str
    uuid: str
    build_metadata: EnvdBuild
